from datasets import load_dataset
from transformers import AutoTokenizer

# 1️⃣ 加载IMDb数据（train / test）
dataset = load_dataset("imdb")

# 2️⃣ 从原始 train 划分验证集（eval），留 test 不动
split_datasets = dataset["train"].train_test_split(test_size=0.1, seed=42)
train_dataset = split_datasets["train"]    # 训练集
eval_dataset  = split_datasets["test"]     # 验证集
test_dataset  = dataset["test"]            # IMDb 原本的测试集

# 3️⃣ 加载分词器
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

# 4️⃣ 定义分词函数
def tokenize_function(examples):
    return tokenizer(
        examples["text"],  #限制必须文本 可以其他的
        padding="max_length",   # 补全到模型最大长度
        truncation=True
    )

# 5️⃣ Tokenization（batched=True 批量加速）
tokenized_train = train_dataset.map(tokenize_function, batched=True)
tokenized_eval  = eval_dataset.map(tokenize_function, batched=True)
tokenized_test  = test_dataset.map(tokenize_function, batched=True)

# 6️⃣ 打印示例
print(tokenized_train[0])
print("---------------------")
print(tokenized_eval[0])